import {z} from 'zod'

import {qwen} from "../models.ts";
import {Annotation, END, Send, START, StateGraph} from "@langchain/langgraph";
import {saveArrayBufferAsImage} from "../tool.mjs";

const subjectsPrompt = "Generate a comma separated list of between 2 and 5 examples related to: {topic}."
const jokePrompt = "Generate a joke about {subject}"
const bestJokePrompt = `Below are a bunch of jokes about {topic}. Select the best one! Return the ID (index) of the best one.

{jokes}`



const Subjects = z.object({
    subjects: z.array(z.string())
});

const Joke = z.object({
    joke: z.string()
})

const BestJoke = z.object({
    id: z.number()
})

const OverallState = Annotation.Root({
    topic: Annotation,
    subjects: Annotation,
    jokes: Annotation({
        reducer: (state, update) => state.concat(update),
    }),
    bestSelectedJoke: Annotation
})
const generateTopics = async (
    state
) => {
    const prompt = subjectsPrompt.replace("topic", state.topic);
    const response = await qwen
        .withStructuredOutput(Subjects, { name: "subjects" })
        .invoke(prompt);
    return { subjects: response.subjects };
};

const generateJoke = async (state) => {
    const prompt = jokePrompt.replace("subject", state.subject);
    const response = await qwen
        .withStructuredOutput(Joke, { name: "joke" })
        .invoke(prompt);
    return { jokes: [response.joke] };
};

// Here we define the logic to map out over the generated subjects
// We will use this an edge in the graph
const continueToJokes = (state) => {
    // We will return a list of `Send` objects
    // Each `Send` object consists of the name of a node in the graph
    // as well as the state to send to that node
    return state.subjects.map((subject) => new Send("generateJoke", { subject }));
};
// Here we will judge the best joke
const bestJoke = async (
    state
)=> {
    const jokes = state.jokes.join("\n\n");
    const prompt = bestJokePrompt
        .replace("jokes", jokes)
        .replace("topic", state.topic);
    const response = await qwen
        .withStructuredOutput(BestJoke, { name: "best_joke" })
        .invoke(prompt);
    return { bestSelectedJoke: state.jokes[response.id] };
};
// Construct the graph: here we put everything together to construct our graph
const graph = new StateGraph(OverallState)
    .addNode("generateTopics", generateTopics)
    .addNode("generateJoke", generateJoke)
    .addNode("bestJoke", bestJoke)
    .addEdge(START, "generateTopics")
    .addConditionalEdges("generateTopics", continueToJokes)
    .addEdge("generateJoke", "bestJoke")
    .addEdge("bestJoke", END);

const app = graph.compile();
const representation = await app.getGraphAsync()

const image = await representation.drawMermaidPng();
const arrayBuffer = await image.arrayBuffer();

saveArrayBufferAsImage(arrayBuffer, "flow.png")


for await (const s of await app.stream({ topic: "animals" })) {
    console.log(s);
}
